#include <iostream>
#include <math.h>

extern "C" {
#include <freerdp/freerdp.h>
#include <freerdp/listener.h>
}
#include <winpr/path.h>
#include <winpr/thread.h>

struct server_info
{
    const char *cert;
    const char *key;
};

int pos_x = 0;
int pos_y = 0;

void send_black_screen(freerdp_peer *client)
{
    const int width = 1920;
    const int height = 1080;
    const int bpp = 32;
    const int stride = width * (bpp / 8); // 每行数据大小
    const int data_size = stride * height;

    // 分配白色屏幕数据
    BYTE *white_data = (BYTE *)calloc(1, data_size);
    if (!white_data) {
        return;
    }

    uint32_t color = 0xff000000;

    // 逐像素填充白色 (ARGB 格式)
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            int index = (y * width + x) * 4;
            uint32_t *pixel = (uint32_t *)(white_data + index);
            if (sqrt((x - pos_x) * (x - pos_x) + (y - pos_y) * (y - pos_y)) < 20) {
                *pixel = 0xffffffff; // red
            }
        }
    }

    SURFACE_BITS_COMMAND cmd = {0};
    cmd.cmdType = CMDTYPE_SET_SURFACE_BITS;
    cmd.destLeft = 0;
    cmd.destTop = 0;
    cmd.destRight = width;
    cmd.destBottom = height;
    cmd.bmp.bpp = 32;
    cmd.bmp.width = width;
    cmd.bmp.height = height;
    cmd.bmp.bitmapData = white_data;
    cmd.bmp.bitmapDataLength = data_size;

    // 发送屏幕数据
    client->context->update->SurfaceBits(client->context, &cmd);

    free(white_data); // 释放内存
}

BOOL handleClientActivate(freerdp_peer *client)
{
    std::cout << "Activate" << std::endl;
    send_black_screen(client);
    return true;
}

BOOL handleClientPostConnect(freerdp_peer *client)
{
    std::cout << "PostConnect" << std::endl;
    send_black_screen(client);
    return true;
}

DWORD WINAPI client_main_loop(freerdp_peer *client)
{
    server_info *info = static_cast<server_info *>(client->ContextExtra);

    if (!freerdp_peer_context_new(client)) {
        freerdp_peer_free(client);
        return -1;
    }

    rdpSettings *settings = client->context->settings;

    rdpPrivateKey *key = freerdp_key_new_from_file(info->key);
    freerdp_settings_set_pointer_len(settings, FreeRDP_RdpServerRsaKey, (void *)key, 1);

    rdpCertificate *cert = freerdp_certificate_new_from_file(info->cert);
    freerdp_settings_set_pointer_len(settings, FreeRDP_RdpServerCertificate, (void *)cert, 1);

    freerdp_settings_set_bool(settings, FreeRDP_RdpSecurity, TRUE);
    freerdp_settings_set_bool(settings, FreeRDP_TlsSecurity, TRUE);
    freerdp_settings_set_bool(settings, FreeRDP_NlaSecurity, FALSE);
    freerdp_settings_set_uint32(settings, FreeRDP_EncryptionLevel, ENCRYPTION_LEVEL_CLIENT_COMPATIBLE);
    freerdp_settings_set_bool(settings, FreeRDP_RemoteFxCodec, TRUE);
    freerdp_settings_set_bool(settings, FreeRDP_NSCodec, TRUE);

    // 动态获取
    freerdp_settings_set_uint32(settings, FreeRDP_DesktopWidth, 1920);
    freerdp_settings_set_uint32(settings, FreeRDP_DesktopHeight, 1080);
    freerdp_settings_set_uint32(settings, FreeRDP_ColorDepth, 32);

    freerdp_settings_set_bool(settings, FreeRDP_SuppressOutput, TRUE);
    freerdp_settings_set_bool(settings, FreeRDP_RefreshRect, TRUE);
    freerdp_settings_set_uint32(settings, FreeRDP_MultifragMaxRequestSize, 0xFFFFFF /* FIXME */);

    // send a blank screen to the client
    client->PostConnect = handleClientPostConnect;
    client->Activate = handleClientActivate;

    // 处理输入
    client->context->input->KeyboardEvent = [](rdpInput *input, UINT16 flags, UINT8 code) -> BOOL {
        return TRUE;
    };

    client->context->input->MouseEvent = [](rdpInput *input, UINT16 flags, UINT16 x, UINT16 y) -> BOOL {
        pos_x = x;
        pos_y = 1080 - y;
        return TRUE;
    };

    while (client->CheckFileDescriptor(client) == TRUE) {
        send_black_screen(client);
        // usleep(10 * 1000);
    }

    std::cout << "Client Main Loop Quit!" << std::endl;

    freerdp_peer_context_free(client);
    freerdp_peer_free(client);
    return 0;
}

BOOL handlePeerAccepted(freerdp_listener *instance, freerdp_peer *client)
{
    std::cout << "Peer accepted from " << client->hostname << std::endl;
    server_info *info = static_cast<server_info *>(instance->info);
    client->ContextExtra = info;

    HANDLE clientThread = nullptr;
    clientThread = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)client_main_loop, client, 0, NULL);
    if (!clientThread) {
        return false;
    }

    CloseHandle(clientThread);
    return true;
}

int main(int argc, char *argv[])
{
    // Initialize FreeRDP
    freerdp_listener *listener = freerdp_listener_new();
    if (!listener) {
        return -1;
    }

    server_info info;
    info.cert = "/etc/krd/tls.crt";
    info.key = "/etc/krd/tls.key";

    listener->info = &info;
    listener->PeerAccepted = handlePeerAccepted;
    listener->Open(listener, NULL, 3389);

    while (1) {
        listener->GetEventHandles(listener, NULL, 0);

        if (listener->CheckFileDescriptor(listener) != true) {
            break;
        }
    }

    return 0;
}